package net.lab1024.sa.admin.module.business.pdf.util;

import lombok.extern.slf4j.Slf4j;
import net.sourceforge.tess4j.Tesseract;
import org.springframework.stereotype.Component;

import jakarta.annotation.Resource;
import java.awt.*;
import java.awt.image.BufferedImage;
import java.awt.image.ConvolveOp;
import java.awt.image.Kernel;
import java.awt.image.RescaleOp;
import java.util.HashMap;
import java.util.Map;

/**
 * 自适应OCR优化器
 * 根据图像质量分析结果自动调整OCR参数和预处理策略
 *
 * @Author 1024创新实验室
 */
@Slf4j
@Component
public class AdaptiveOcrOptimizer {

    @Resource
    private OcrDiagnosticTool diagnosticTool;

    /**
     * 根据图像质量分析结果生成优化策略
     */
    public OcrOptimizationStrategy generateOptimizationStrategy(BufferedImage image) {
        Map<String, Object> imageAnalysis = diagnosticTool.analyzeImageQuality(image);
        OcrOptimizationStrategy strategy = new OcrOptimizationStrategy();
        
        // 解析图像质量数据
        double avgBrightness = (Double) imageAnalysis.get("averageBrightness");
        double contrast = (Double) imageAnalysis.get("contrast");
        int width = (Integer) imageAnalysis.get("width");
        int height = (Integer) imageAnalysis.get("height");
        String qualityLevel = (String) imageAnalysis.get("qualityLevel");
        
        log.info("图像质量分析: 亮度={}, 对比度={}, 尺寸={}x{}, 质量等级={}", 
                avgBrightness, contrast, width, height, qualityLevel);
        
        // 根据质量等级确定基础策略
        switch (qualityLevel) {
            case "优秀":
                strategy = generateExcellentStrategy(avgBrightness, contrast, width, height);
                break;
            case "良好":
                strategy = generateGoodStrategy(avgBrightness, contrast, width, height);
                break;
            case "一般":
                strategy = generateFairStrategy(avgBrightness, contrast, width, height);
                break;
            case "较差":
                strategy = generatePoorStrategy(avgBrightness, contrast, width, height);
                break;
        }
        
        // 根据具体参数进行微调
        strategy = fineTuneStrategy(strategy, avgBrightness, contrast, width, height);
        
        log.info("生成的优化策略: {}", strategy);
        return strategy;
    }
    
    /**
     * 优秀质量图像的策略
     */
    private OcrOptimizationStrategy generateExcellentStrategy(double brightness, double contrast, int width, int height) {
        OcrOptimizationStrategy strategy = new OcrOptimizationStrategy();
        
        // 优秀图像使用轻微优化
        strategy.setPreprocessingLevel("minimal");
        strategy.setContrastEnhancement(1.05f);
        strategy.setBrightnessAdjustment(0.0f);
        strategy.setSharpenLevel(0);
        strategy.setBinarizationThreshold(0.7f);
        strategy.setUseGrayscale(true);
        strategy.setUseBinary(false);
        
        // Tesseract参数优化
        strategy.setPageSegMode(3); // 自动页面分割
        strategy.setOcrEngineMode(3); // LSTM引擎
        strategy.setMinLineSize(0.8f);
        strategy.setNoiseRemoval(false);
        strategy.setPreserveInterwordSpaces(true);
        
        return strategy;
    }
    
    /**
     * 良好质量图像的策略
     */
    private OcrOptimizationStrategy generateGoodStrategy(double brightness, double contrast, int width, int height) {
        OcrOptimizationStrategy strategy = new OcrOptimizationStrategy();
        
        // 良好图像使用中等优化
        strategy.setPreprocessingLevel("moderate");
        strategy.setContrastEnhancement(1.15f);
        strategy.setBrightnessAdjustment(calculateBrightnessAdjustment(brightness));
        strategy.setSharpenLevel(1);
        strategy.setBinarizationThreshold(0.6f);
        strategy.setUseGrayscale(true);
        strategy.setUseBinary(false);
        
        // Tesseract参数
        strategy.setPageSegMode(3);
        strategy.setOcrEngineMode(3);
        strategy.setMinLineSize(0.6f);
        strategy.setNoiseRemoval(false);
        strategy.setPreserveInterwordSpaces(true);
        
        return strategy;
    }
    
    /**
     * 一般质量图像的策略
     */
    private OcrOptimizationStrategy generateFairStrategy(double brightness, double contrast, int width, int height) {
        OcrOptimizationStrategy strategy = new OcrOptimizationStrategy();
        
        // 一般图像使用较强优化
        strategy.setPreprocessingLevel("aggressive");
        strategy.setContrastEnhancement(1.3f);
        strategy.setBrightnessAdjustment(calculateBrightnessAdjustment(brightness));
        strategy.setSharpenLevel(2);
        strategy.setBinarizationThreshold(0.5f);
        strategy.setUseGrayscale(true);
        strategy.setUseBinary(true); // 尝试二值化
        
        // Tesseract参数
        strategy.setPageSegMode(6); // 统一文本块
        strategy.setOcrEngineMode(3);
        strategy.setMinLineSize(0.4f);
        strategy.setNoiseRemoval(true);
        strategy.setPreserveInterwordSpaces(true);
        
        return strategy;
    }
    
    /**
     * 较差质量图像的策略
     */
    private OcrOptimizationStrategy generatePoorStrategy(double brightness, double contrast, int width, int height) {
        OcrOptimizationStrategy strategy = new OcrOptimizationStrategy();
        
        // 较差图像使用最强优化
        strategy.setPreprocessingLevel("extreme");
        strategy.setContrastEnhancement(1.5f);
        strategy.setBrightnessAdjustment(calculateBrightnessAdjustment(brightness));
        strategy.setSharpenLevel(3);
        strategy.setBinarizationThreshold(0.4f);
        strategy.setUseGrayscale(true);
        strategy.setUseBinary(true);
        
        // Tesseract参数
        strategy.setPageSegMode(7); // 单行文本
        strategy.setOcrEngineMode(3);
        strategy.setMinLineSize(0.2f);
        strategy.setNoiseRemoval(true);
        strategy.setPreserveInterwordSpaces(false);
        
        return strategy;
    }
    
    /**
     * 根据亮度计算亮度调整值
     */
    private float calculateBrightnessAdjustment(double brightness) {
        if (brightness < 80) {
            return 30.0f; // 过暗，增加亮度
        } else if (brightness > 200) {
            return -30.0f; // 过亮，降低亮度
        } else {
            return 0.0f; // 正常亮度
        }
    }
    
    /**
     * 微调策略
     */
    private OcrOptimizationStrategy fineTuneStrategy(OcrOptimizationStrategy strategy, 
                                                   double brightness, double contrast, int width, int height) {
        
        // 根据对比度微调
        if (contrast < 0.3) {
            strategy.setContrastEnhancement(strategy.getContrastEnhancement() * 1.2f);
        } else if (contrast > 0.8) {
            strategy.setContrastEnhancement(strategy.getContrastEnhancement() * 0.9f);
        }
        
        // 根据分辨率微调
        int minDimension = Math.min(width, height);
        if (minDimension < 800) {
            strategy.setMinLineSize(strategy.getMinLineSize() * 0.8f);
        } else if (minDimension > 2000) {
            strategy.setMinLineSize(strategy.getMinLineSize() * 1.2f);
        }
        
        return strategy;
    }
    
    /**
     * 应用优化策略到图像
     */
    public BufferedImage applyOptimizationStrategy(BufferedImage originalImage, OcrOptimizationStrategy strategy) {
        BufferedImage processedImage = originalImage;
        
        try {
            // 1. 亮度调整
            if (strategy.getBrightnessAdjustment() != 0) {
                processedImage = adjustBrightness(processedImage, strategy.getBrightnessAdjustment());
            }
            
            // 2. 对比度增强
            if (strategy.getContrastEnhancement() != 1.0f) {
                processedImage = enhanceContrast(processedImage, strategy.getContrastEnhancement());
            }
            
            // 3. 锐化
            if (strategy.getSharpenLevel() > 0) {
                processedImage = sharpenImage(processedImage, strategy.getSharpenLevel());
            }
            
            // 4. 转换为灰度
            if (strategy.isUseGrayscale()) {
                processedImage = convertToGrayscale(processedImage);
            }
            
            // 5. 二值化
            if (strategy.isUseBinary()) {
                processedImage = convertToBinary(processedImage, strategy.getBinarizationThreshold());
            }
            
            log.info("图像预处理完成，策略: {}", strategy.getPreprocessingLevel());
            
        } catch (Exception e) {
            log.warn("图像预处理失败，使用原图: {}", e.getMessage());
            return originalImage;
        }
        
        return processedImage;
    }
    
    /**
     * 应用优化策略到Tesseract实例
     */
    public void applyOptimizationStrategy(Tesseract tesseract, OcrOptimizationStrategy strategy) {
        // 页面分割模式
        tesseract.setTessVariable("tessedit_pageseg_mode", String.valueOf(strategy.getPageSegMode()));
        
        // OCR引擎模式
        tesseract.setTessVariable("tessedit_ocr_engine_mode", String.valueOf(strategy.getOcrEngineMode()));
        
        // 最小行大小
        tesseract.setTessVariable("textord_min_linesize", String.valueOf(strategy.getMinLineSize()));
        
        // 噪声移除
        tesseract.setTessVariable("textord_heavy_nr", strategy.isNoiseRemoval() ? "1" : "0");
        
        // 保留词间空格
        tesseract.setTessVariable("preserve_interword_spaces", strategy.isPreserveInterwordSpaces() ? "1" : "0");
        
        // 中文识别优化参数
        tesseract.setTessVariable("language_model_penalty_non_dict_word", "0.8");
        tesseract.setTessVariable("language_model_penalty_non_freq_dict_word", "0.8");
        tesseract.setTessVariable("segment_penalty_dict_nonword", "0.8");
        
        log.info("Tesseract参数已应用，策略: {}", strategy.getPreprocessingLevel());
    }
    
    // 图像处理辅助方法
    private BufferedImage adjustBrightness(BufferedImage image, float adjustment) {
        RescaleOp rescale = new RescaleOp(1.0f, adjustment, null);
        return rescale.filter(image, null);
    }
    
    private BufferedImage enhanceContrast(BufferedImage image, float scaleFactor) {
        RescaleOp rescale = new RescaleOp(scaleFactor, 0.0f, null);
        return rescale.filter(image, null);
    }
    
    private BufferedImage sharpenImage(BufferedImage image, int level) {
        float[] kernel;
        switch (level) {
            case 1: // 轻微锐化
                kernel = new float[]{
                    0.0f, -0.25f, 0.0f,
                    -0.25f, 2.0f, -0.25f,
                    0.0f, -0.25f, 0.0f
                };
                break;
            case 2: // 中等锐化
                kernel = new float[]{
                    0.0f, -0.5f, 0.0f,
                    -0.5f, 3.0f, -0.5f,
                    0.0f, -0.5f, 0.0f
                };
                break;
            case 3: // 强烈锐化
                kernel = new float[]{
                    0.0f, -1.0f, 0.0f,
                    -1.0f, 5.0f, -1.0f,
                    0.0f, -1.0f, 0.0f
                };
                break;
            default:
                return image;
        }
        
        ConvolveOp convolve = new ConvolveOp(new Kernel(3, 3, kernel), ConvolveOp.EDGE_NO_OP, null);
        return convolve.filter(image, null);
    }
    
    private BufferedImage convertToGrayscale(BufferedImage image) {
        BufferedImage grayImage = new BufferedImage(image.getWidth(), image.getHeight(), BufferedImage.TYPE_BYTE_GRAY);
        Graphics2D g = grayImage.createGraphics();
        g.drawImage(image, 0, 0, null);
        g.dispose();
        return grayImage;
    }
    
    private BufferedImage convertToBinary(BufferedImage image, float threshold) {
        BufferedImage grayImage = convertToGrayscale(image);
        BufferedImage binaryImage = new BufferedImage(grayImage.getWidth(), grayImage.getHeight(), BufferedImage.TYPE_BYTE_BINARY);
        
        int thresholdValue = (int) (threshold * 255);
        
        for (int y = 0; y < grayImage.getHeight(); y++) {
            for (int x = 0; x < grayImage.getWidth(); x++) {
                int rgb = grayImage.getRGB(x, y);
                int gray = rgb & 0xFF;
                int binary = gray > thresholdValue ? 0xFFFFFF : 0x000000;
                binaryImage.setRGB(x, y, binary);
            }
        }
        
        return binaryImage;
    }
    
    /**
     * OCR优化策略类
     */
    public static class OcrOptimizationStrategy {
        private String preprocessingLevel = "minimal";
        private float contrastEnhancement = 1.0f;
        private float brightnessAdjustment = 0.0f;
        private int sharpenLevel = 0;
        private float binarizationThreshold = 0.7f;
        private boolean useGrayscale = true;
        private boolean useBinary = false;
        
        // Tesseract参数
        private int pageSegMode = 3;
        private int ocrEngineMode = 3;
        private float minLineSize = 0.8f;
        private boolean noiseRemoval = false;
        private boolean preserveInterwordSpaces = true;
        
        // Getters and Setters
        public String getPreprocessingLevel() { return preprocessingLevel; }
        public void setPreprocessingLevel(String preprocessingLevel) { this.preprocessingLevel = preprocessingLevel; }
        
        public float getContrastEnhancement() { return contrastEnhancement; }
        public void setContrastEnhancement(float contrastEnhancement) { this.contrastEnhancement = contrastEnhancement; }
        
        public float getBrightnessAdjustment() { return brightnessAdjustment; }
        public void setBrightnessAdjustment(float brightnessAdjustment) { this.brightnessAdjustment = brightnessAdjustment; }
        
        public int getSharpenLevel() { return sharpenLevel; }
        public void setSharpenLevel(int sharpenLevel) { this.sharpenLevel = sharpenLevel; }
        
        public float getBinarizationThreshold() { return binarizationThreshold; }
        public void setBinarizationThreshold(float binarizationThreshold) { this.binarizationThreshold = binarizationThreshold; }
        
        public boolean isUseGrayscale() { return useGrayscale; }
        public void setUseGrayscale(boolean useGrayscale) { this.useGrayscale = useGrayscale; }
        
        public boolean isUseBinary() { return useBinary; }
        public void setUseBinary(boolean useBinary) { this.useBinary = useBinary; }
        
        public int getPageSegMode() { return pageSegMode; }
        public void setPageSegMode(int pageSegMode) { this.pageSegMode = pageSegMode; }
        
        public int getOcrEngineMode() { return ocrEngineMode; }
        public void setOcrEngineMode(int ocrEngineMode) { this.ocrEngineMode = ocrEngineMode; }
        
        public float getMinLineSize() { return minLineSize; }
        public void setMinLineSize(float minLineSize) { this.minLineSize = minLineSize; }
        
        public boolean isNoiseRemoval() { return noiseRemoval; }
        public void setNoiseRemoval(boolean noiseRemoval) { this.noiseRemoval = noiseRemoval; }
        
        public boolean isPreserveInterwordSpaces() { return preserveInterwordSpaces; }
        public void setPreserveInterwordSpaces(boolean preserveInterwordSpaces) { this.preserveInterwordSpaces = preserveInterwordSpaces; }
        
        @Override
        public String toString() {
            return String.format("OcrOptimizationStrategy{preprocessingLevel='%s', contrastEnhancement=%.2f, " +
                    "brightnessAdjustment=%.1f, sharpenLevel=%d, binarizationThreshold=%.2f, " +
                    "useGrayscale=%s, useBinary=%s, pageSegMode=%d, minLineSize=%.2f, noiseRemoval=%s}",
                    preprocessingLevel, contrastEnhancement, brightnessAdjustment, sharpenLevel, 
                    binarizationThreshold, useGrayscale, useBinary, pageSegMode, minLineSize, noiseRemoval);
        }
    }
} 